package com.nanohadoop.hdfs;

import com.nanohadoop.utils.Configuration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class DataNode {
    private static final Logger log = LoggerFactory.getLogger(DataNode.class);

    private final String id;
    private final String host;
    private final int port;
    private final int ipcPort;
    private final Configuration conf;
    private final File dataDir;
    private final ScheduledExecutorService scheduler;
    private final Map<Long, File> blockFiles;
    private final List<Long> blocks;
    
    private ServerSocket serverSocket;
    private Thread serverThread;
    private volatile boolean running;

    public DataNode(String id) {
        this(id, Configuration.getInstance());
    }

    public DataNode(String id, Configuration conf) {
        this.conf = conf;
        this.host = conf.get("dfs.datanode.hostname", "localhost");
        this.port = conf.getInt("dfs.datanode.port", 50010);
        this.ipcPort = conf.getInt("dfs.datanode.ipc.port", 50020);
        this.id = host + ":" + port;  // ID format: host:port
        
        // 初始化数据目录
        String dataDirPath = conf.get("dfs.datanode.data.dir", "file:///tmp/dfs/data");
        dataDirPath = dataDirPath.replace("file://", "");
        this.dataDir = new File(dataDirPath, "datanode-" + id);
        initializeDataDirectory();

        this.blockFiles = new HashMap<>();
        this.blocks = new ArrayList<>();
        this.scheduler = Executors.newScheduledThreadPool(1);

        // 启动心跳和块报告任务
        startHeartbeatTask();
    }

    private void initializeDataDirectory() {
        if (!dataDir.exists()) {
            if (!dataDir.mkdirs()) {
                throw new RuntimeException("Failed to create data directory: " + dataDir);
            }
        }
    }

    private void startHeartbeatTask() {
        // 每30秒发送一次心跳
        scheduler.scheduleAtFixedRate(this::sendHeartbeat, 0, 30, TimeUnit.SECONDS);
        // 每小时发送一次块报告
        scheduler.scheduleAtFixedRate(this::sendBlockReport, 0, 1, TimeUnit.HOURS);
    }

    public void sendHeartbeat() {
        try (Socket socket = new Socket(conf.get("fs.defaultFS.host", "localhost"),
                                      conf.getInt("dfs.namenode.port", 9000));
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
                
            out.writeUTF("HEARTBEAT");
            out.writeUTF(id);
            // 发送存储信息
            long totalSpace = dataDir.getTotalSpace();
            long usedSpace = totalSpace - dataDir.getFreeSpace();
            out.writeLong(totalSpace);
            out.writeLong(usedSpace);
            // 发送块列表
            updateBlockList();
            out.writeInt(blocks.size());
            for (Long blockId : blocks) {
                out.writeLong(blockId);
            }
            out.flush();

            // 读取响应
            boolean success = in.readBoolean();
            if (success) {
                log.debug("Heartbeat sent successfully from DataNode: " + id);
            }
        } catch (Exception e) {
            System.err.println("Failed to send heartbeat: " + e.getMessage());
        }
    }

    public void sendBlockReport() {
        try {
            // TODO: 实现发送块报告到NameNode的逻辑
            // 扫描数据目录获取所有块
            updateBlockList();
            // 发送块列表
        } catch (Exception e) {
            System.err.println("Failed to send block report: " + e.getMessage());
        }
    }

    private void updateBlockList() {
        blocks.clear();
        File[] files = dataDir.listFiles((dir, name) -> name.startsWith("blk_"));
        if (files != null) {
            for (File file : files) {
                try {
                    long blockId = Long.parseLong(file.getName().replace("blk_", ""));
                    blocks.add(blockId);
                    blockFiles.put(blockId, file);
                } catch (NumberFormatException e) {
                    // 忽略无效的块文件名
                }
            }
        }
    }

    public void start() throws IOException {
        if (running) {
            return;
        }
        running = true;

        // 启动服务器
        serverSocket = new ServerSocket(port);
        serverThread = new Thread(() -> {
            while (running) {
                try {
                    Socket socket = serverSocket.accept();
                    // 处理客户端请求
                    new Thread(() -> handleClientRequest(socket)).start();
                } catch (IOException e) {
                    if (running) {
                        e.printStackTrace();
                    }
                }
            }
        });
        serverThread.start();

        // 发送注册信息到 NameNode
        registerWithNameNode();
    }

    public void stop() {
        if (!running) {
            return;
        }
        running = false;
        
        try {
            if (serverSocket != null) {
                serverSocket.close();
            }
            if (serverThread != null) {
                serverThread.join();
            }
        } catch (IOException | InterruptedException e) {
            e.printStackTrace();
        }
        
        shutdown();
    }

    private void handleClientRequest(Socket socket) {
        try (DataInputStream in = new DataInputStream(socket.getInputStream());
             DataOutputStream out = new DataOutputStream(socket.getOutputStream())) {
            
            String command = in.readUTF();
            log.debug("DataNode received command: " + command);
            
            switch (command) {
                case "READ":
                    long blockId = in.readLong();
                    log.debug("DataNode " + id + " reading block: " + blockId);
                    byte[] data = readBlock(blockId);
                    log.debug("DataNode " + id + " read " + data.length + " bytes from block " + blockId);
                    out.writeInt(data.length);
                    out.write(data);
                    out.flush();
                    break;
                    
                case "WRITE":
                    blockId = in.readLong();
                    int length = in.readInt();
                    log.debug("DataNode " + id + " receiving block: " + blockId + ", length: " + length);
                    data = new byte[length];
                    in.readFully(data);
                    try {
                        writeBlock(blockId, data);
                        out.writeBoolean(true);  // 写入成功
                        out.flush();
                        log.debug("DataNode " + id + ": Block " + blockId + " written successfully");
                    } catch (IOException e) {
                        out.writeBoolean(false);  // 写入失败
                        out.flush();
                        System.err.println("DataNode " + id + ": Failed to write block " + blockId + ": " + e.getMessage());
                        throw e;
                    }
                    break;
                    
                default:
                    System.err.println("Unknown command: " + command);
                    out.writeBoolean(false);
            }
            
            out.flush();
            
        } catch (IOException e) {
            System.err.println("Error handling client request: " + e.getMessage());
            e.printStackTrace();
        } finally {
            try {
                socket.close();
            } catch (IOException e) {
                System.err.println("Error closing socket: " + e.getMessage());
            }
        }
    }

    private void registerWithNameNode() throws IOException {
        log.debug("Registering DataNode with NameNode...");
        Socket socket = new Socket(conf.get("fs.defaultFS.host", "localhost"),
                                 conf.getInt("dfs.namenode.port", 9000));
        try (DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            out.writeUTF("REGISTER_DATANODE");
            String nodeId = host + ":" + port;
            out.writeUTF(nodeId);  // 发送格式化的ID (host:port)
            out.writeUTF(host);
            out.writeInt(port);
            out.flush();
            
            boolean success = in.readBoolean();
            if (success) {
                log.debug("DataNode registered successfully with ID: " + nodeId);
            } else {
                throw new IOException("Failed to register DataNode");
            }
        } catch (IOException e) {
            System.err.println("Failed to register DataNode: " + e.getMessage());
            throw e;
        } finally {
            socket.close();
        }
    }

    public void writeBlock(long blockId, byte[] data) throws IOException {
        File blockFile = new File(dataDir, "blk_" + blockId);
        
        // 确保目录存在
        blockFile.getParentFile().mkdirs();
        
        log.debug("Writing block " + blockId + " to " + blockFile.getAbsolutePath());
        try (FileOutputStream fos = new FileOutputStream(blockFile)) {
            fos.write(data);
            fos.flush();
        } catch (IOException e) {
            System.err.println("Failed to write block file: " + e.getMessage());
            if (blockFile.exists()) {
                blockFile.delete();
            }
            throw e;
        }
        
        blockFiles.put(blockId, blockFile);
        blocks.add(blockId);
        log.debug("Successfully wrote block " + blockId + ", size: " + data.length + " bytes");
    }

    public byte[] readBlock(long blockId) throws IOException {
        // 先尝试从缓存中获取
        File blockFile = blockFiles.get(blockId);
        if (blockFile == null || !blockFile.exists()) {
            // 如果缓存中没有，则根据命名规则查找
            blockFile = new File(dataDir, "blk_" + blockId);
            if (!blockFile.exists()) {
                throw new IOException("Block " + blockId + " not found in DataNode " + id);
            }
            // 更新缓存
            blockFiles.put(blockId, blockFile);
        }
        
        log.debug("DataNode " + id + " reading from file: " + blockFile.getAbsolutePath());
        byte[] data = new byte[(int) blockFile.length()];
        try (FileInputStream fis = new FileInputStream(blockFile)) {
            int bytesRead = fis.read(data);
            log.debug("DataNode " + id + " read " + bytesRead + " bytes from block " + blockId);
        }
        return data;
    }

    public String getId() {
        return id;
    }

    public String getHost() {
        return host;
    }

    public int getPort() {
        return port;
    }

    public int getIpcPort() {
        return ipcPort;
    }

    public List<Long> getBlocks() {
        return new ArrayList<>(blocks);
    }

    public void shutdown() {
        scheduler.shutdown();
        try {
            scheduler.awaitTermination(5, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
}